import requests
import json
from flask import Flask, request

from requests import Timeout

app = Flask(__name__)

headers = {
    'accept': 'application/json',
    'Content-Type': 'application/json',
}


def chat(host="http://192.168.1.4",
         port="7861",
         api="/chat/chat",
         query="你好！",
         history=[],
         stream=False,
         top_k=3,
         score_threshold=1.0,
         temperature=0.7,
         max_tokens=None
         ):
    url = f"{host}:{port}{api}"
    data = {
        "query": query,
        "history": history,
        "stream": stream,
        "top_k": top_k,
        "score_threshold": score_threshold,
        "temperature": temperature,
        "max_tokens": max_tokens
    }
    # 设置超时时间，例如总超时时间为10秒
    timeout = 30
    result = []
    print(f"query：{query}")
    try:
        response = requests.post(url, headers=headers, json=data, stream=stream, timeout=timeout)
        # 处理响应
        if response.status_code == 200:
            # 处理成功的响应
            for line in response.iter_content(None, decode_unicode=True):
                data = json.loads(line[6:])
                result.append(data)
            res = {
                "msg": "请求成功",
                "code": 200,
                "data": result
            }
            print(f"请求成功，返回：{result[0]['text']}")
        else:
            # 处理失败的响应
            print(f"请求失败，状态码：{response.status_code}")
            res = {
                "msg": "请求失败",
                "code": response.status_code,
                "data": response.text
            }
        return res
    except Timeout as e:
        # 捕获超时异常
        print(f"请求超时：{e}")
        # 在这里处理超时的情况，比如重试请求、记录日志或向用户显示错误消息
        res = {
            "msg": "请求失败",
            "code": 504,
            "data": e
        }
        return res
    except requests.exceptions.RequestException as e:
        # 捕获其他类型的请求异常
        print(f"请求发生错误：{e}")
        # 在这里处理其他类型的请求异常
        res = {
            "msg": "请求失败",
            "code": 504,
            "data": e
        }
        return res


# 和大模型对话获取文字描述
@app.route('/chat_to_llm', methods=['POST'])
def chat_to_llm():
    query = request.form['query']
    return chat(query=query)


if __name__ == "__main__":
    app.config['JSON_AS_ASCII'] = False
    # 将host设置为0.0.0.0，则外网用户也可以访问到这个服务
    app.run(host="0.0.0.0", port=5333, debug=False)
